import os
import random
from random import shuffle
import numpy as np
import torch
from torch.utils import data
from torchvision import transforms as T
from torchvision.transforms import functional as F
from PIL import Image
import cv2
import numpy
from albumentations import (
    Compose,
    GaussianBlur,
    HorizontalFlip,
    MedianBlur,
    MotionBlur,
    Normalize,
    OneOf,
    RandomBrightness,
    RandomContrast,
    Resize,
    ShiftScaleRotate,
    VerticalFlip,
    CLAHE,
	Rotate,
	GaussNoise,
GridDistortion,ElasticTransform
)

class ImageFolder(data.Dataset):
	def __init__(self, root,image_size=512,mode='train',augmentation_prob=0.4):
		"""Initializes image paths and preprocessing module."""
		self.root = root
		
		# GT : Ground Truth
		self.GT_paths = root[:-1]+'_GT/'
		self.image_paths = list(map(lambda x: os.path.join(root, x), os.listdir(root)))
		self.image_size = image_size
		self.mode = mode
		self.RotationDegree = [0,90,180,270]
		self.augmentation_prob = augmentation_prob
		self.transform=generate_transforms((384,384))
		print("image count in {} path :{}".format(self.mode,len(self.image_paths)))

	def __getitem__(self, index):
		"""Reads an image from a file and preprocesses it and returns."""
		image_path = self.image_paths[index]
		filename = image_path.split('/')[-1]
		GT_path = self.GT_paths  + filename.replace(".jpg"," mask.png")
		'''image = np.array(Image.open(image_path).convert('RGB'))
		GT = np.array(Image.open(GT_path).convert('L'))'''
		image = cv2.cvtColor(
			cv2.imread(image_path), cv2.COLOR_BGR2RGB
		)
		GT = cv2.cvtColor(
			cv2.imread(GT_path), cv2.COLOR_BGR2GRAY
		)/255


		if (self.mode == 'train'):
			transformed = self.transform['train_transforms'](image=image, mask=GT)
			image = transformed["image"]
			GT = transformed["mask"]
			transformed = self.transform['color_transforms'](image=image)
			image = transformed["image"]
			transformed = self.transform['val_transforms'](image=image)
			image = transformed["image"]
		else:
			transformed = self.transform['val_transforms'](image=image)
			image = transformed["image"]


		transformed = self.transform['resize'](image=image,mask=GT)
		image = transformed["image"]
		GT = transformed["mask"]
		image = image.transpose(2, 0, 1)


		return image, GT,filename

	def __len__(self):
		"""Returns the total number of font files."""
		return len(self.image_paths)

def generate_transforms(image_size):

	train_transform = Compose(
        [VerticalFlip(p=0.5),
		 Rotate(limit=45, interpolation=1, border_mode=0, value=(0,0,0), mask_value=(0,0,0), always_apply=False, p=0.5),
		 GridDistortion(num_steps=5, distort_limit=0.3, interpolation=1, border_mode=0, value=(0,0,0),
						mask_value=(0,0,0),
						always_apply=False, p=0.2 ),
		 ElasticTransform(alpha = 1, sigma = 50, alpha_affine = 50, interpolation = 1,
						  border_mode = 0, value =(255,255,255), mask_value=(0,0,0),
						  always_apply = False, approximate = False, p = 0.2 ),
		 HorizontalFlip(p=0.5),])
	color_transform = Compose(
		[CLAHE(clip_limit=4.0, tile_grid_size=(8, 8), always_apply=False, p=0.0),
		GaussNoise(var_limit=(10.0, 50.0), mean=0, always_apply=False, p=0.0), ])

	val_transform = Compose(
        [
		 Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5),max_pixel_value=255.0, p=1.0),])

	resize_transform = Compose(
		[Resize(height=image_size[0], width=image_size[1]), ])

	return {"train_transforms": train_transform, "val_transforms": val_transform,"resize": resize_transform,"color_transforms":color_transform}



def get_loader(image_path, image_size, batch_size, num_workers=2, mode='train',augmentation_prob=0.4):
	"""Builds and returns Dataloader."""
	
	dataset = ImageFolder(root = image_path, image_size =image_size, mode=mode,augmentation_prob=augmentation_prob)
	data_loader = data.DataLoader(dataset=dataset,
								  batch_size=batch_size,
								  shuffle=True,
								  num_workers=num_workers)
	return data_loader

def get_loader_GAN(image_path, image_size, batch_size, num_workers=2, mode='train',augmentation_prob=0.4):
	"""Builds and returns Dataloader."""
	
	dataset = ImageFolder_GAN(root = image_path, image_size =image_size, mode=mode,augmentation_prob=augmentation_prob)
	data_loader = data.DataLoader(dataset=dataset,
								  batch_size=batch_size,
								  shuffle=True,
								  num_workers=num_workers)
	return data_loader


class ImageFolder_GAN(data.Dataset):
	def __init__(self, root,image_size=512,mode='train',augmentation_prob=0.4):
		"""Initializes image paths and preprocessing module."""
		self.root = root
		
		# GT : Ground Truth
		self.GT_paths = root[:-1]+'_GT/'
		self.image_paths = list(map(lambda x: os.path.join(root, x), os.listdir(root)))
		self.image_paths_GAN = list(map(lambda x: os.path.join(root[:-1]+'_GAN/', x), os.listdir(root[:-1]+'_GAN/')))
		self.image_size = image_size
		self.mode = mode
		self.RotationDegree = [0,90,180,270]
		self.augmentation_prob = augmentation_prob
		self.transform=generate_transforms((384,384))
		print("image count in {} path :{}".format(self.mode,len(self.image_paths)))

	def __getitem__(self, index):
		"""Reads an image from a file and preprocesses it and returns."""
		image_path = self.image_paths[index]
		image_path_GAN = self.image_paths_GAN[index%len(self.image_paths_GAN)]
		filename = image_path.split('/')[-1]
		GT_path = self.GT_paths + filename.replace(".jpg"," mask.png")
		'''image = np.array(Image.open(image_path).convert('RGB'))
		GT = np.array(Image.open(GT_path).convert('L'))'''
		image = cv2.cvtColor(
			cv2.imread(image_path), cv2.COLOR_BGR2RGB
		)
		image_GAN = cv2.cvtColor(
			cv2.imread(image_path_GAN), cv2.COLOR_BGR2RGB
		)
		GT = cv2.cvtColor(
			cv2.imread(GT_path), cv2.COLOR_BGR2GRAY
		)/255


		if (self.mode == 'train'):
			transformed = self.transform['train_transforms'](image=image, mask=GT)
			image = transformed["image"]
			GT = transformed["mask"]
			#transformed = self.transform['train_transforms'](image=image_GAN)
			#image_GAN = transformed["image"]

			transformed = self.transform['color_transforms'](image=image)
			image = transformed["image"]
			#transformed = self.transform['color_transforms'](image=image_GAN)
			#image_GAN = transformed["image"]

			transformed = self.transform['val_transforms'](image=image)
			image = transformed["image"]
			transformed = self.transform['val_transforms'](image=image_GAN)
			image_GAN = transformed["image"]
		else:
			transformed = self.transform['val_transforms'](image=image)
			image = transformed["image"]


		transformed = self.transform['resize'](image=image,mask=GT)
		image = transformed["image"]
		GT = transformed["mask"]

		transformed = self.transform['resize'](image=image_GAN)
		image_GAN = transformed["image"]

		image = image.transpose(2, 0, 1)
		image_GAN = image_GAN.transpose(2, 0, 1)

		return image, GT,filename,image_GAN

	def __len__(self):
		"""Returns the total number of font files."""
		return len(self.image_paths)